package xyz.thoughtset.viewer.ai.mcp.server.factory;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import xyz.thoughtset.viewer.executor.blocks.tool.BlockToolCallbackProvider;
import xyz.thoughtset.viewer.ai.mcp.server.config.DynamicRouterRegistry;
import xyz.thoughtset.viewer.ai.mcp.server.entity.McpServerInfo;
import xyz.thoughtset.viewer.ai.mcp.server.exception.UrlConflictException;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
public class McpServerFactory implements DisposableBean {
    private final Map<String, McpServerInfo> MCP_SERVER_INFO_MAP = new ConcurrentHashMap<>();
    private final Map<String, McpSyncServer> MCP_SERVER_MAP = new ConcurrentHashMap<>();
    private final Set<String> MCP_SSE_SET = new CopyOnWriteArraySet<>();
    private final Set<String> MCP_MSG_SET = new CopyOnWriteArraySet<>();
    @Autowired
    private DynamicRouterRegistry dynamicRouterRegistry;
    @Autowired
    private ObjectMapper objectMapper;
    @Autowired
    private BlockToolCallbackProvider blockToolCallbackProvider;

    @Override
    public void destroy() {
        MCP_SERVER_MAP.values().parallelStream().forEach(McpSyncServer::close);
    }

    public void createMcpServer(List<McpServerInfo> list) throws UrlConflictException {
        if (ObjectUtils.isEmpty(list)) return;
        HashSet<String> sseSet = new HashSet<>();
        HashSet<String> msgSet = new HashSet<>();
        Iterator<McpServerInfo> iterator = list.iterator();
        while (iterator.hasNext()) {
            McpServerInfo ele = iterator.next();
            String sse = ele.getSseEndpoint();
            String msg = ele.getMessageEndpoint();
            if (sseSet.contains(sse) || msgSet.contains(msg)) {
                throw UrlConflictException.build();
            }
            createMcpServer(ele);
            sseSet.add(sse);
            msgSet.add(msg);
        }
        MCP_SSE_SET.addAll(sseSet);
        MCP_MSG_SET.addAll(msgSet);
    }

    public void createMcpServer(McpServerInfo mcpServerInfo) throws UrlConflictException {
        String key = mcpServerInfo.getId();
        String baseUrl = mcpServerInfo.getBaseUrl();
        String sseEndpoint = mcpServerInfo.getSseEndpoint();
        String messageEndpoint = mcpServerInfo.getMessageEndpoint();
        if (!(MCP_SSE_SET.isEmpty() && MCP_MSG_SET.isEmpty()) &&
                (MCP_SSE_SET.contains(sseEndpoint) || MCP_MSG_SET.contains(messageEndpoint))) {
            throw UrlConflictException.build();
        }
        WebMvcSseServerTransportProvider.Builder transportProviderBuilder = WebMvcSseServerTransportProvider.builder()
//                .jsonMapper(new JacksonMcpJsonMapper(objectMapper))
                .objectMapper(objectMapper)
                .baseUrl(baseUrl)
                .sseEndpoint(sseEndpoint);
        if (StringUtils.hasText(messageEndpoint)){
            transportProviderBuilder.messageEndpoint(messageEndpoint);
        }
        Integer interval = mcpServerInfo.getKeepAliveInterval();
        if (interval != null && interval.longValue() > 0) {
            transportProviderBuilder.keepAliveInterval(Duration.ofSeconds(interval));
        }
        WebMvcSseServerTransportProvider transportProvider = transportProviderBuilder.build();
        dynamicRouterRegistry.register(key, transportProvider);
        McpSchema.ServerCapabilities.Builder capabilityBuilder = McpSchema.ServerCapabilities.builder().logging();
//        var capabilities = McpSchema.ServerCapabilities.builder()
//                .logging() // Logging support
//                .build();
        McpServer.SyncSpecification serverBuilder = McpServer.sync(transportProvider)
                .serverInfo(mcpServerInfo.getTitle(), mcpServerInfo.getVersion());
//                .serverInfo("MCP Demo Weather Server", "1.0.0")
        mgtTools(capabilityBuilder,serverBuilder,key);
        //promptly, we disable file and function support
        serverBuilder.instructions(mcpServerInfo.getInstructions());
        Integer timeout = mcpServerInfo.getRequestTimeout();
        if (timeout != null && timeout.longValue() > 0) {
            serverBuilder.requestTimeout(Duration.ofSeconds(timeout));
        }
//        serverBuilder.immediateExecution(mcpServerInfo.isImmediateExecution());
        serverBuilder.capabilities(capabilityBuilder.build());
        McpSyncServer server = serverBuilder.build();
        MCP_SERVER_MAP.put(key, server);
        MCP_SERVER_INFO_MAP.put(key, mcpServerInfo);
    }

    private void mgtTools(McpSchema.ServerCapabilities.Builder capabilityBuilder ,
                          McpServer.SyncSpecification serverBuilder ,
                          String serverId) {
        capabilityBuilder.tools(true); // Tool support with list changes notifications
        serverBuilder.tools(McpToolUtils.toSyncToolSpecifications(
                blockToolCallbackProvider.getToolCallbacks(serverId)
        ));
    }

    public void closeMcpServer(String key) {
        McpSyncServer server = MCP_SERVER_MAP.remove(key);
        dynamicRouterRegistry.unregister(key);
        if (server != null) {
            server.close();
        }
        McpServerInfo serverInfo = MCP_SERVER_INFO_MAP.remove(key);
        MCP_SSE_SET.remove(serverInfo.getSseEndpoint());
        MCP_MSG_SET.remove(serverInfo.getMessageEndpoint());
    }

}
